from abc import abstractclassmethod, abstractmethod, abstractproperty
from typing import Self
import struct
import math
import unittest
import hashlib
import hmac
import functools
from Crypto.Cipher import ARC4 as O_ARC4
import random
from os import urandom

class hashcls:
    def __init__(self):
        self.buf = b''

    @classmethod
    def new(cls, msg: bytes = b''):
        c = cls()
        c.update(msg)
        return c

    def update(self: Self, data: bytes):
        self.buf += data

    def hexdigest(self: Self):
        return ''.join([f'{x:02x}' for x in self.digest()])

    @property
    def block_size(self: Self):
        return 64

    @staticmethod
    def ROL(v: int, c: int):
        v &= 0xFFFFFFFF
        return ((v << c) | v >> (32 - c)) & 0xFFFFFFFF

    @abstractmethod
    def digest(self: Self)->bytes:
        pass

    @abstractproperty
    def digest_size(self: Self)->int:
        pass

class MD5(hashcls):
    def __init__(self: Self):
        super(MD5, self).__init__()

    @property
    def digest_size(self: Self)->int:
        return 16

    def digest(self: Self):
        buf = self.buf + b'\x80' + (119- len(self.buf)) % 64 * b'\x00' + struct.pack('<Q', len(self.buf) * 8)
        A, = struct.unpack('<I', struct.pack('>I', 0x01234567)) # 0x67452301
        B, = struct.unpack('<I', struct.pack('>I', 0x89ABCDEF)) # 0xEFCDAB89
        C, = struct.unpack('<I', struct.pack('>I', 0xFEDCBA98)) # 0x98BADCFE
        D, = struct.unpack('<I', struct.pack('>I', 0x76543210)) # 0x10325476

        F = lambda x, y, z: (x & y) | (~x & z)
        G = lambda x, y, z: (x & z) | (y & ~z)
        H = lambda x, y, z: x ^ y ^ z
        I = lambda x, y, z: y ^ (x | ~z)
        T = lambda i: math.floor(2**32 * abs(math.sin(i)))

        SS = [
            [7, 12, 17, 22],
            [5,  9, 14, 20],
            [4, 11, 16, 23],
            [6, 10, 15, 21],
        ]

        for i in range(0, len(buf), 64):
            X = [struct.unpack('<I', buf[i+j : i+j+4])[0] for j in range(0, 64, 4)]
            AA = A
            BB = B
            CC = C
            DD = D

            Z = lambda a,b,c,d,k,s,i,fn: (b + hashcls.ROL(a + fn(b,c,d) + X[k] + T(i), s)) & 0xFFFFFFFF

            for j in range(0, 64):
                A, D, C, B = D, C, B, Z(A, B, C, D, (j * [1, 5, 3, 7][j // 16] + [0, 1, 5, 0][j // 16]) % 16, SS[j // 16][j % 4], j+1, [F, G, H, I][j // 16])

            A = (A + AA) & 0xFFFFFFFF
            B = (B + BB) & 0xFFFFFFFF
            C = (C + CC) & 0xFFFFFFFF
            D = (D + DD) & 0xFFFFFFFF

        return struct.pack('<4I', A, B, C, D)

class SHA1(hashcls):
    def __init__(self: Self):
        super(SHA1, self).__init__()

    @property
    def digest_size(self: Self)->int:
        return 20

    def digest(self: Self):
        buf = self.buf + b'\x80' + (119- len(self.buf) % 64) % 64 * b'\x00' + struct.pack('>Q', len(self.buf) * 8)
        A, = struct.unpack('<I', struct.pack('>I', 0x01234567)) # 0x67452301
        B, = struct.unpack('<I', struct.pack('>I', 0x89ABCDEF)) # 0xEFCDAB89
        C, = struct.unpack('<I', struct.pack('>I', 0xFEDCBA98)) # 0x98BADCFE
        D, = struct.unpack('<I', struct.pack('>I', 0x76543210)) # 0x10325476
        E, = struct.unpack('<I', struct.pack('>I', 0xF0E1D2C3)) # 0xC3D2E1F0

        for i in range(0, len(buf), 64):
            w = []
            for j in range(0, 64, 4):
                w.append(struct.unpack('>I', buf[i+j : i+j+4])[0])
            for j in range(16, 80):
                w.append(hashcls.ROL(w[j-3]^w[j-8]^w[j-14]^w[j-16], 1))
            AA = A
            BB = B
            CC = C
            DD = D
            EE = E

            funcs = [
                lambda b,c,d: ((b & c) | ((~b + 0x100000000) & d), 0x5A827999),
                lambda b,c,d: (b ^ c ^ d, 0x6ED9EBA1),
                lambda b,c,d: ((b & c) | (b & d) | (c & d), 0x8F1BBCDC),
                lambda b,c,d: (b ^ c ^ d, 0xCA62C1D6),
            ]
            for j in range(80):
                f,k = funcs[j // 20](B,C,D)
                A,B,C,D,E = (hashcls.ROL(A, 5) + f + E + k + w[j]) & 0xFFFFFFFF, A, hashcls.ROL(B, 30), C, D
            A = (A + AA) & 0xFFFFFFFF
            B = (B + BB) & 0xFFFFFFFF
            C = (C + CC) & 0xFFFFFFFF
            D = (D + DD) & 0xFFFFFFFF
            E = (E + EE) & 0xFFFFFFFF
        return struct.pack('>5I', A, B, C, D, E)

def HMAC(key: bytes, msg: bytes, digestmod: hashcls):
    outer = digestmod.new()
    inner = digestmod.new()
    blocksize = inner.block_size

    if len(key) > blocksize:
        key = digestmod.new(key).digest()

    key += b'\0' * (blocksize - len(key))
    outer.update(bytes([x ^ 0x5C for x in key]))
    inner.update(bytes([x ^ 0x36 for x in key]))
    if msg is not None:
        inner.update(msg)

    outer.update(inner.digest())
    return outer.digest()

class DerInteger:
    def __init__(self, data: bytes):
        self.__raw_data = data
        self.__parse()

    @property
    def raw(self):
        return self.__raw_data

    @property
    def size_in_bytes(self):
        return len(self.__raw_data.lstrip(b'\x00'))

    @property
    def value(self):
        return self.__value

    def __parse(self):
        if (self.__raw_data[0] & 0x80) == 0:
            self.__value = functools.reduce(lambda x,y:x<<8|y, self.__raw_data, 0)
        else:
            self.__value = functools.reduce(lambda x,y:x<<8|(0xFF^y), self.__raw_data, 0)
            self.__value = -self.__value-1


def der_decode(data: bytes):
    def TLV(data: bytes, check_length: bool = True):
        t,l = data[:2]
        i = 2
        if (l & 0x80) > 0:
            ll = l & 0x7F
            l = functools.reduce(lambda x,y:x<<8|y,data[2:2+ll],0)
            i += ll
        if check_length:
            assert len(data) == i + l
            return t,l,data[i:]
        else:
            return t,l,data[:i+l],data[i+l:]

    t,l,v = TLV(data)
    match t:
        case 1: # BOOLEAN
            assert l == 1
            v = v[0]
            assert v in (0, 0xFF)
            return v == 0xFF
        case 2: # INTEGER
            return DerInteger(v)
        case 3: # BIT STRING
            return v[0], v[1:]
        case 4: # OCTET STRING
            return v
        case 5: # NULL
            assert l == 0
            return None
        case 6: # OBJECT_ID
            t = []
            t.append(v[0] // 40)
            t.append(v[0] % 40)
            v = v[1:]
            while len(v) > 0:
                iii = 0
                z = v[0] & 0x7F
                while v[iii] >= 0x80:
                    z <<= 7
                    iii += 1
                    z |= v[iii] & 0x7F
                v = v[iii+1:]
                t.append(z)
            return tuple(t)
        case 0xC: # UTF8String
            return v.decode('utf8')
        case 0x13 | 0x16 | 0x17: # PrintableString
            return v.decode()
        case 0x14: # TeletexString
            # print('TeletexString')
            return ''
        case 0x1E: # BMPString
            return v.decode('utf-16be')
        case 0x30 | 0xA0 | 0xA3:
            r = []
            while len(v) > 0:
                t,l,vv,v=TLV(v, False)
                r.append(vv)
            return r
        case 0x31:
            r = set()
            while len(v) > 0:
                t,l,v=TLV(v, False)
                r.add(v[:l])
                v = v[l:]
            return r
        case _:
            raise NotImplementedError(f'unknown tag: {t: 02X}, 0x{l:X}: {v[:20]}')

class ARC4:
    def __init__(self, key: bytes):
        self.decrypt = self.encrypt

        self.S = list(range(256))

        j = 0
        for i in range(256):
            j = (j + self.S[i] + key[i % len(key)]) % 256
            self.S[i], self.S[j] = self.S[j], self.S[i]

        self.i = 0
        self.j = 0

    def encrypt(self, value: bytes):
        r = []
        for x in value:
            self.i = (self.i + 1) % 256
            self.j = (self.j + self.S[self.i]) % 256
            self.S[self.i], self.S[self.j] = self.S[self.j], self.S[self.i]
            t = (self.S[self.i] + self.S[self.j]) % 256
            r.append(self.S[t] ^ x)
        return bytes(r)
    
class RSA:
    def __init__(self, n: DerInteger, e: DerInteger):
        self.n = n
        self.e = e

    @classmethod
    def from_der(cls, data):
        '''
        SEQUENCE {
            SEQUENCE {
                OBJECT_ID algoid;
                NULL null; // opt
            };
            BIT STRING {
                Integer module;
                Integer public exponent;
            }
        }
        '''
        d = der_decode(data)
        assert len(d) == 2
        s = der_decode(d[0])
        assert len(s) in [1, 2]
        ss = der_decode(s[0])
        assert isinstance(ss, tuple)
        assert ss == (1,2,840,113549,1,1,1)
        if len(s) == 2:
            ss = der_decode(s[1])
            assert ss is None
        x,k = der_decode(d[1])
        assert x == 0
        k = der_decode(k)
        assert len(k) == 2
        n = der_decode(k[0])
        e = der_decode(k[1])
        assert isinstance(n, DerInteger)
        assert isinstance(e, DerInteger)
        return cls(n, e)
    
    def encrypt(self, message: bytes):
        k = self.n.size_in_bytes
        mLen = len(message)

        # Step 1
        if mLen > k - 11:
            raise ValueError("Plaintext is too long.")
        # Step 2a
        ps = bytes([random.randint(1,255) for _ in range(k - mLen - 3)])
        assert(len(ps) == k - mLen - 3)
        # Step 2b
        em = b'\x00\x02' + ps + b'\x00' + message
        # Step 3a (OS2IP)
        em_int = int.from_bytes(em, 'big')
        # Step 3b (RSAEP)
        m_int = int(pow(em_int, self.e.value, self.n.value))
        # Step 3c (I2OSP)
        c = RSA.long_to_bytes(m_int, k)
        return c
    
    @staticmethod
    def long_to_bytes(n, blocksize=0):
        if n < 0 or blocksize < 0:
            raise ValueError("Values must be non-negative")

        result = []
        pack = struct.pack

        # Fill the first block independently from the value of n
        bsr = blocksize
        while bsr >= 8:
            result.insert(0, pack('>Q', n & 0xFFFFFFFFFFFFFFFF))
            n = n >> 64
            bsr -= 8

        while bsr >= 4:
            result.insert(0, pack('>I', n & 0xFFFFFFFF))
            n = n >> 32
            bsr -= 4

        while bsr > 0:
            result.insert(0, pack('>B', n & 0xFF))
            n = n >> 8
            bsr -= 1

        if n == 0:
            if len(result) == 0:
                bresult = b'\x00'
            else:
                bresult = b''.join(result)
        else:
            # The encoded number exceeds the block size
            while n > 0:
                result.insert(0, pack('>Q', n & 0xFFFFFFFFFFFFFFFF))
                n = n >> 64
            result[0] = result[0].lstrip(b'\x00')
            bresult = b''.join(result)
            # bresult has minimum length here
            if blocksize > 0:
                target_len = ((len(bresult) - 1) // blocksize + 1) * blocksize
                bresult = b'\x00' * (target_len - len(bresult)) + bresult

        return bresult


class TTT(unittest.TestCase):
    def test_hash(self: Self):
        ks = [
            "",
            "a",
            "abc",
            "message digest",
            "abcdefghijklmnopqrstuvwxyz",
            "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789",
            "12345678901234567890123456789012345678901234567890123456789012345678901234567890"
        ]
        for k in ks:
            v = hashlib.md5(k.encode()).hexdigest()
            print(f'{k} md5: {v}')
            m = MD5()
            m.update(k.encode())
            r = m.hexdigest()
            with self.subTest(k=k,v=v):
                self.assertEqual(r, v)

            v = hashlib.sha1(k.encode()).hexdigest()
            print(f'{k} sha1: {v}')
            m = SHA1()
            m.update(k.encode())
            r = m.hexdigest()
            with self.subTest(k=k,v=v):
                self.assertEqual(r, v)

    def test_hmac(self):
        args = [
            (b'', b''),
            (b'1', b'2'),
            (b'0123ABC789', b'0QWE456789'),
            (b'0123ABC789' * 10, b'0QWE456789' * 10),
            (b'0123ABC789' * 30, b'0QWE456789' * 30),
        ]

        for k,m in args:
            print(f'{k} {m} md5')
            with self.subTest(k=k,m=m):
                x = HMAC(k,m,MD5)
                y = hmac.HMAC(k, m, hashlib.md5).digest()
                self.assertEqual(x, y)

            print(f'{k} {m} sha1')
            with self.subTest(k=k,m=m):
                x = HMAC(k,m,SHA1)
                y = hmac.HMAC(k, m, hashlib.sha1).digest()
                self.assertEqual(x, y)

    def test_arc4(self):
        ks = [
            b"",
            b"a",
            b"abc",
            b"message digest",
            b"abcdefghijklmnopqrstuvwxyz",
            b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789",
            b"12345678901234567890123456789012345678901234567890123456789012345678901234567890"
        ]

        key = random.randbytes(16)

        nn = ARC4(key)
        oo = O_ARC4.new(key)
        nnd = ARC4(key)
        ood = O_ARC4.new(key)
        for k in ks:
            if len(k) < 10:
                print(f'{k}')
            else:
                print(f'{k[:10]}...({len(k)})')
            with self.subTest(k=k):
                x=nn.encrypt(k)
                y=oo.encrypt(k)
                self.assertEqual(x, y)
                xx=nnd.decrypt(x)
                yy=ood.decrypt(y)
                self.assertEqual(xx, yy)
                self.assertEqual(xx, k)

if __name__=='__main__':
    TTT().test_hash()
    TTT().test_hmac()
    TTT().test_arc4()